# -*- coding: utf-8 -*-

"""
kg json to Gremlin TinkerPOP GraphSON format

Reference the following url for format detail: https://developer.ibm.com/dwblog/2017/tips-help-graphson-tinkerpop/
"""

from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division

import logging

from kgpipeline.job import KgJob
from kgpipeline.jsonutil import entity_json_decode, entity_json_encode
from kgpipeline.sparkutil import to_text_dataframe

logger = logging.getLevelName(__name__)


def vertex_id(item):
    entity, index = item
    return {"id": index, "label": entity.get("@type", [""])[0], "properties": entity}


def link_out_v_id(item):
    link, vertex = item
    link["outV"] = vertex["id"]
    return link


def link_in_v_id(item):
    link, vertex = item
    link["inV"] = vertex["id"]
    return link


def merge_out_e(item):
    vertex, edge = item
    if edge is not None:
        vertex["outE"] = edge
    return vertex


def merge_in_e(item):
    vertex, edge = item
    if edge is not None:
        vertex["inE"] = edge
    return vertex


class KgGraphSONJob(KgJob):
    """job for convert json to graphson"""

    def process(self, inputs):
        logger.info("Start KgGraphSONJob: {} ...".format(self.name))

        if len(inputs) != 2:
            logger.error("This job takes exactly 2 text input (json) sources... abort!")
            return []

        rdd_list = []

        for df, _ in inputs:
            rdd_list.append(df.rdd.map(lambda x: entity_json_decode(x["value"])))

        entity, link = rdd_list

        res = []

        if not entity.isEmpty():
            vertices = entity.zipWithIndex().map(vertex_id)

            if not link.isEmpty():
                vertices_with_kgid = vertices.map(lambda x: (x["properties"]["@id"], x))
                vertices_with_id = vertices.map(lambda x: (x["id"], x))
                link_start = vertices.count()
                link_in = link.zipWithIndex().map(lambda x: (x[0]["in"], {"id": x[1] + link_start, "properties": x[0]}))
                links = link_in.leftOuterJoin(vertices_with_kgid).mapValues(link_out_v_id).map(
                    lambda x: (x[1]["properties"]["out"], x[1])
                ).leftOuterJoin(vertices_with_kgid).mapValues(link_in_v_id).map(lambda x: x[1])

                vertices_out_e = links.map(lambda x: ((x["outV"], x["properties"]["@type"][0]), x)).groupByKey().map(
                    lambda x: (x[0][0], (x[0][1], x[1]))
                ).groupByKey().mapValues(dict)

                vertices_in_e = links.map(lambda x: ((x["inV"], x["properties"]["@type"][0]), x)).groupByKey().map(
                    lambda x: (x[0][0], (x[0][1], x[1]))
                ).groupByKey().mapValues(dict)

                vertices_with_id = vertices_with_id.leftOuterJoin(vertices_out_e).mapValues(merge_out_e)
                vertices_with_id = vertices_with_id.leftOuterJoin(vertices_in_e).mapValues(merge_in_e)

                vertices = vertices_with_id.map(lambda x: x[1])

            res.append(to_text_dataframe(vertices, entity_json_encode))

        return res


KgJob.register("graphson", KgGraphSONJob)
